import os
import json
import random
from typing import List, Dict, Optional

import datasets
from langchain.prompts import PromptTemplate
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field
import concurrent.futures
from tqdm import tqdm

intention = 'CLAIM_PROCESS'

# 定义结构化输出模型
class QuestionAnswer(BaseModel):
    answer: str = Field(description="模拟真实的客户发言文本")


# 提示模板定义
prompt_template = PromptTemplate(
    input_variables=["question"],
    template="Answer the following question:\n\nQuestion: {question}\n\nAnswer:"
)


def get_structured_output(model_class, prompt: PromptTemplate, question: str) -> BaseModel:
    """
    使用 LLM 获取结构化输出。
    """
    os.environ["OPENAI_API_KEY"] = 'local-qwen2.5-72b-little-brother'
    base_url = "http://10.8.50.33:8814/v1"
    model_name = "qwq-32b-preview"

    llm = ChatOpenAI(
        base_url=base_url,
        model=model_name,
        api_key=os.environ["OPENAI_API_KEY"],
        max_tokens=1000
    )

    chain = prompt | llm.with_structured_output(model_class)
    return chain.invoke({"question": question})


def load_all_few_shots(csv_path: str, label_filter: str = "意识到对方为AI") -> List[str]:
    """
    加载所有符合条件的 few-shot 样本。

    Args:
        csv_path (str): CSV 文件路径。
        label_filter (str): 筛选 label。

    Returns:
        List[str]: 所有符合条件的 prompt 列表。
    """
    dataset = datasets.load_dataset('csv', data_files=csv_path)['train']
    filtered_samples = dataset.filter(lambda x: x['answer'] == label_filter)
    return filtered_samples['prompt']


def load_knowledge_data(knowledge_path: str = '/Users/sd/Desktop/mycode/myalgo/milvus/data-gen/zero_gen/knowledge.json') -> Dict:
    """
    加载知识数据。

    Args:
        knowledge_path (str): 知识文件路径。

    Returns:
        Dict: 知识数据字典。
    """
    with open(knowledge_path, 'r', encoding='utf-8') as f:
        return json.load(f)


def build_prompt(few_shots: List[str]) -> str:
    """
    构建用于生成客户质疑销售为机器人的对话提示语，包含 few-shot 示例。

    Args:
        few_shots (List[str]): few-shot 示例列表。

    Returns:
        str: 构建好的提示字符串。
    """
    # 加载知识数据
    knowledge_data = load_knowledge_data()
    
    # 随机选择一些知识内容添加到提示中
    added_knowledge = []
    
    # 从每个字段类型中随机选择1-2个条目
    for category, items in knowledge_data.items():
        # 随机选择1-2个条目
        num_to_select = random.randint(1, 2) if len(items) >= 2 else 1
        selected_items = random.sample(items, min(num_to_select, len(items)))
        added_knowledge.extend([f"{category}: {item}" for item in selected_items])
    
    # 随机打乱添加的知识内容顺序
    random.shuffle(added_knowledge)
    
    # 组装 few-shot 示例为带案例标题的格式
    examples_str = ""
    for i, ex in enumerate(few_shots, 1):
        examples_str += f"## 案例 {i}\n{ex}\n\n"

    intention_descriptions = {
        "PRODUCT_QA": "关于保险产品具体问题的询问，如产品保障范围、条款细节、投保规则等",
        "PRODUCT_ALL_CONTENT": "询问保险产品的完整保障内容，如保障责任、除外责任等",
        "PRODUCT_REC": "请求推荐合适的保险产品，如根据年龄、健康状况等条件推荐产品",
        "OTHER": "其他类型的问题或陈述，如基本咨询、一般性问题等",
        "APPOINTMENT_CONSULTATION": "预约人工咨询服务，如要求转接人工客服",
        "PRODUCT_HB": "关于保险产品条款的询问，如保险责任、免责条款等",
        "PRODUCT_HEALTH": "关于健康保险的询问，如健康告知、核保规则等",
        "PRODUCT_HB_REC": "请求推荐意外险产品",
        "PRODUCT_FEE": "关于保险费用的询问，如保费计算、费率表等",
        "PRODUCT_BUY": "关于购买保险流程的询问，如投保链接、购买方式等",
        "PRODUCT_COMPARE": "比较不同保险产品，如优缺点对比、性价比分析等",
        "PREMIUM_CALCULATION": "保险费用计算相关问题，如具体保费金额查询",
        "INSURANCE_COMPANY": "关于保险公司的问题，如公司背景、偿付能力等",
        "PRODUCT_PLAN": "关于保险计划的问题，如家庭保险规划、个人保障方案等",
        "INSURANCE_KNOWLEDGE": "保险知识相关问题，如保险术语解释、保险原理等",
        "PRODUCT_DUTY": "关于保险责任的问题，如具体保障范围、赔付条件等",
        "BENEFIT_ILLUSTRATION": "关于保险利益说明的问题，如现金价值、年金领取等",
        "CLAIM_PROCESS": "关于问题咨询多围绕理赔流程,后续理赔规则,理赔步骤等",
        "APPLICATION_RULES": "关于产品怎么投保、投保需要什么条件、投保年龄限制、在哪能买等"
    }



    knowledge_section = "\n".join(added_knowledge) if added_knowledge else ""

    return f"""
# 任务目标：
生成单句提问，客户在了解相关保险过程中，提问或者相关咨询，想了解说的话，提问1个问题。

# 用户的意图
{intention_descriptions[intention]}

# 可参考的背景信息（随机选取）：
{knowledge_section}

# 输出格式：
- 返回 JSON 格式，字段包括 [answer]
- [answer]字段中只包含客户提问句子，每条用换行分隔


# Few-Shot 示例：
{examples_str.strip()}
## 最终案例：

"""


def analyze_customer_talk(dialog: str, all_few_shots: List[str], sample_size: int = 10) -> Dict:
    """
    对单条输入调用大模型生成客户质疑为机器人的话语，并注入随机 few-shot 示例。

    Args:
        dialog (str): 输入内容（仅占位）
        all_few_shots (List[str]): 所有可用的 few-shot 样本
        sample_size (int): 每次采样的数量

    Returns:
        dict: 包含生成结果的字典。
    """
    try:
        # ✅ 每次调用都独立随机选取 10 个样本
        few_shots = random.sample(all_few_shots, min(sample_size, len(all_few_shots)))
        response = get_structured_output(QuestionAnswer, prompt_template, build_prompt(few_shots))
        return response.model_dump()
    except Exception as e:
        print(f"Error analyzing: '{dialog}' - {e}")
        return {"item": dialog, "error": str(e)}


def process_item(item: str, all_few_shots: List[str], sample_size: int = 10) -> Dict:
    """
    单条数据处理函数，用于并发调用生成，每次使用不同随机 few-shot。

    Args:
        item (str): 输入占位符
        all_few_shots (List[str]): 所有可用的 few-shot 样本
        sample_size (int): 每次采样的数量

    Returns:
        dict: 包含生成结果的字典。
    """
    return analyze_customer_talk(item, all_few_shots, sample_size)


def generate_multiple_samples(
        num_samples: int = 50,
        save_path: Optional[str] = None,
        csv_path: str = '/mnt/data2/zcz/infer/utils/notebook/data/unlike_label2_ai_tag_v2.csv',
        label_filter: str = "意识到对方为AI",
        sample_size: int = 10
) -> List[Dict]:
    """
    并发生成多个客户质疑为机器人的对话样本，每次使用不同随机 few-shot。

    Args:
        num_samples (int): 需要生成的样本数量。
        save_path (Optional[str]): 中间结果保存路径。
        csv_path (str): few-shot 数据来源路径。
        label_filter (str): 筛选 label。
        sample_size (int): 每次采样的 few-shot 数量。

    Returns:
        List[Dict]: 所有生成的对话列表。
    """
    # ✅ 一次性加载所有 few-shot 样本
    all_few_shots = load_all_few_shots(csv_path, label_filter)
    print(f"Loaded {len(all_few_shots)} total few-shot samples.")

    results = []
    cnt = 0

    with concurrent.futures.ThreadPoolExecutor(max_workers=12) as executor:
        futures = {executor.submit(process_item, "", all_few_shots, sample_size): i for i in range(num_samples)}

        for future in tqdm(concurrent.futures.as_completed(futures), total=len(futures), colour="GREEN"):
            try:
                result = future.result(timeout=600)
                results.append(result)

                # 每处理 50 条自动保存一次
                if cnt % 50 == 0 and save_path:
                    with open(save_path, 'w', encoding='utf-8') as f:
                        json.dump(results, f, indent=2, ensure_ascii=False)

                cnt += 1

            except Exception as e:
                print(f"Task failed with exception: {e}")
                cnt += 1

    # 最终保存一次
    if save_path:
        with open(save_path, 'w', encoding='utf-8') as f:
            json.dump(results, f, indent=2, ensure_ascii=False)

    return results


if __name__ == "__main__":
    # 设置参数
    output_path = '/Users/sd/Desktop/mycode/myalgo/milvus/data-gen/zero_gen/generated_normal2.json'
    csv_path = '/Users/sd/Desktop/mycode/myalgo/milvus/data-gen/zero_gen/intention.csv'

    # 生成 50 条样本并保存
    generated_data = generate_multiple_samples(
        num_samples=5000,
        save_path=output_path,
        csv_path=csv_path,
        label_filter=intention,  # "意识到对方为AI",
        sample_size=5  # 每次调用都随机抽取 10 个示例
    )

    # 打印部分结果
    print(json.dumps(generated_data[:], ensure_ascii=False, indent=4))
    print("生成完成！")